import os
import hashlib
import argparse


def calculate_md5(file_path, block_size=65536):
    """计算单个文件的MD5值"""
    md5_hash = hashlib.md5()
    try:
        with open(file_path, "rb") as f:
            # 分块读取大文件，避免内存占用过高
            for block in iter(lambda: f.read(block_size), b""):
                md5_hash.update(block)
        return md5_hash.hexdigest()
    except Exception as e:
        print(f"计算文件 {file_path} 的MD5时出错: {str(e)}")
        return None


def batch_calculate_md5(root_dir, output_file="md5_results.txt", include_subdirs=True):
    """批量计算目录下所有文件的MD5值"""
    results = []

    # 遍历目录
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for filename in filenames:
            file_path = os.path.join(dirpath, filename)
            # 计算相对路径，使结果更清晰
            relative_path = os.path.relpath(file_path, root_dir)
            print(f"正在处理: {relative_path}")

            md5_value = calculate_md5(file_path)
            if md5_value:
                results.append((relative_path, md5_value))

    # 保存结果到文件
    with open(output_file, "w", encoding="utf-8") as f:
        for path, md5 in results:
            f.write(f"{md5}  {path}\n")

    print(f"\n处理完成，共计算了 {len(results)} 个文件的MD5值")
    print(f"结果已保存到: {os.path.abspath(output_file)}")


if __name__ == "__main__":
    # 设置命令行参数
    parser = argparse.ArgumentParser(description="批量计算文件的MD5值")
    parser.add_argument("directory", help="要处理的目录路径")
    parser.add_argument("-o", "--output", help="结果输出文件路径", default="md5_results.txt")
    parser.add_argument("-ns", "--no-subdirs", action="store_true",
                        help="不包含子目录中的文件")

    args = parser.parse_args()

    # 验证目录是否存在
    if not os.path.isdir(args.directory):
        print(f"错误: 目录 '{args.directory}' 不存在")
        exit(1)

    # 调用批量计算函数
    batch_calculate_md5(
        root_dir=args.directory,
        output_file=args.output,
        include_subdirs=not args.no_subdirs
    )
